# 要添加一个新单元，输入 '# %%'
# 要添加一个新的标记单元，输入 '# %% [markdown]'
# 我发现写到最后，还是学弟写的那种样子。好麻烦啊，感觉应该一天就能写完，但自己在硬拖。
import copy
import torch
from torch.utils import data
import os
import numpy as np
import logging
from trainers.fedbase.FedClientBase import FedClientBase

# logging.basicConfig(level = logging.INFO,format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logging.basicConfig(format ='%(message)s')
logger = logging.getLogger()
logger.setLevel(logging.INFO)

class FedAvgClient(FedClientBase):
    """训练相关的参数本来应该是服务器发送过来的，包括模型、参数啥的。训练数据应该是本地的。\n
    这里为了方便，只模拟联邦学习的效果和逻辑，不模拟联邦学习的真是过程。

    1. 数据由服务器分成独立同分布，或者非独立同分布。然后分配给各个客户端。而非客户端在本地读取。
    2. 训练的模型、参数在客户端配置。
    3. 客户端完成训练过程、服务端完成聚合过程。客户端和服务端只发送和传输模型和优化器的权重参数。完成聚合和分配的通信逻辑。
    接收参数，完成训练。主要是一次训练的参数。

    每一个算法都应该提供梯度的返回值。即在客户端计算出梯度。其实在服务器计算也无可厚非。
    Args:
        object ([type]): [description]
    """    

    def __init__(self,datapair,model):
        """创建一个客户端对象。初始化客户端的本地参数。作为基类。初始化所有的基类参数，与模型无关的参数。
        1. device 设备，检查GPU是否可用
        2. user_id 随机生成的客户端标识
        3. model_path 用来存储和加载模型的路径
        4. data_pair x,y本地数据集，ndarray格式
        5. model 模型
        6. optimizer 优化器

        Args:
            datapair (tuple(ndarray,ndarray)): input,output包括输入输出数据
            model (torch.model): 与客户端一致的模型
        """        
        super(FedAvgClient, self).__init__(datapair,model)
        logger.info("FedAvgClient init----")

    def set_parameters(self,client_params):
        """从服务器加载训练的参数。\n
        1. epochs 数据集循环的次数
        2. batch_size 每个batch的大小
        3. inner_lr 客户端优化器的步长
        4. outer_lr 服务器端优化器的步长
        5. update_steps 本地更新的步骤。如果定义了这个变量，则batch_size无效。batch_size=len(dataset)/update_steps
        Args:
            client_params (dict): 客户端进行训练的相关参数
        """        
        super().set_parameters(client_params)
        logger.info("FedAvgClient set_parameters----")

    def train(self):
        """用来训练当前的本地模型。每次训练返回损失值running_loss和running_accuracy

        Returns:
            [(running_loss,running_accuracy)]: 训练损失和训练的准确率
        """     

        self.model.train()
        loss_sum = 0.
        loss_times = 0.
        correct_sum = 0.
        total_size = 0
        for epoch in range(self.epochs):
            for batch_x,batch_y in self.train_dl:
                # 正向传播，计算准确度
                outputs = self.model(batch_x)
                _,predicted = torch.max(outputs,1)
                correct_sum += (predicted == batch_y).sum().item()
                total_size += len(batch_y)
                # 计算loss
                loss = self.loss_func(outputs,batch_y.long())
                loss_sum +=loss.item()
                loss_times+=1
                # 反向传播
                self.optimizer.zero_grad()
                loss.backward()
                # 梯度下降
                self.optimizer.step()

        # print(self.optimizer.state_dict())
        return loss_sum/loss_times,100*correct_sum/total_size

    def test(self):
        """用来测试训练的结果返回准确率

        Returns:
            float: 模型在本地数据上进行测试的准确率
        """        
        self.model.eval()
        correct = 0
        total =0
        with torch.no_grad():
            for features,labels in self.train_dl:
                outputs = self.model(features)
                # print(outputs.shape)
                _,predicted = torch.max(outputs,1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        accuracy = 100*correct/total
        return accuracy